package com.zoe.onelink.logging.trace.handler;

import brave.handler.MutableSpan;
import brave.handler.SpanHandler;
import brave.propagation.TraceContext;
import cn.hutool.core.exceptions.ExceptionUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.zoe.onelink.common.constant.TraceSpanConstant;
import com.zoe.onelink.common.enumeration.OnelinkExceptionCode;
import com.zoe.onelink.common.exception.ApiException;
import com.zoe.onelink.util.OnelinkContextHolder;
import org.slf4j.MDC;

import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * <p>标题: </p>
 * <p>描述: </p>
 * <p>版权: Copyright (c) 2021</p>
 * <p>公司: 智业软件股份有限公司</p>
 *
 * @version: 1.0
 * @author: ruansheng
 * @date: 2021-08-06
 */
public class SpanFinishedHandler extends SpanHandler {

    private static final Set<String> IGNORE_KEYS;
    private static final String AUTH_OPERATOR = "operator";
    private static final String AUTH_NAME = "name";
    private static final String EXTRA_POD_IP;
    private static final String EXTRA_NODE_IP;

    static {
        IGNORE_KEYS = Stream.of(
                "traceId", "spanId", "spanExportable", "X-B3-TraceId", "X-B3-SpanId", "X-Span-Export"
        ).collect(Collectors.toSet());
        EXTRA_POD_IP = Optional.ofNullable(System.getenv("TRACING_POD_IP")).orElse("");
        EXTRA_NODE_IP = Optional.ofNullable(System.getenv("TRACING_NODE_IP")).orElse("");
    }

    @Override
    public boolean end(TraceContext context, MutableSpan span, Cause cause) {
        this.appendMdcInfo(span);
        this.fillAuthInfo(span);
        this.countContentLength(span);
        this.appendError(span);
        this.appendExtra(span);
        return true;
    }

    private void appendExtra(MutableSpan span) {
        span.tag(TraceSpanConstant.EXTRA_POD_IP, EXTRA_POD_IP);
        span.tag(TraceSpanConstant.EXTRA_NODE_IP, EXTRA_NODE_IP);
    }

    private void fillAuthInfo(MutableSpan span) {
        String authInfoTag = span.tag(TraceSpanConstant.AUTH_INFO);
        if (StrUtil.isNotBlank(authInfoTag)) {
            return;
        }
        String operator = OnelinkContextHolder.getString(AUTH_OPERATOR);
        if (StrUtil.isNotBlank(operator)) {
            JSONObject authInfo = new JSONObject()
                    .fluentPut(AUTH_OPERATOR, operator)
                    .fluentPut(AUTH_NAME, StrUtil.nullToEmpty(OnelinkContextHolder.getString(AUTH_NAME)));
            span.tag(TraceSpanConstant.AUTH_INFO, authInfo.toJSONString());
        }
    }

    private void appendError(MutableSpan span) {
        Throwable ex = span.error();
        if (ex == null) {
            span.tag(TraceSpanConstant.LOG_LEVEL, "INFO");
            return;
        } else if (ex instanceof ApiException) {
            span.tag(TraceSpanConstant.LOG_LEVEL, "WARN");
        } else {
            span.tag(TraceSpanConstant.LOG_LEVEL, "ERROR");
        }
        String exGroup = ex instanceof ApiException && !(((ApiException) ex).getExceptionCode() instanceof OnelinkExceptionCode)
                ? TraceSpanConstant.ExGroup.BIZ : TraceSpanConstant.ExGroup.SYSTEM;
        span.tag(TraceSpanConstant.EXCEPTION_GROUP, exGroup);
        span.tag(TraceSpanConstant.EXCEPTION_CLASS_NAME, StrUtil.nullToEmpty(ex.getClass().getName()));
        span.tag(TraceSpanConstant.EXCEPTION_MESSAGE, StrUtil.nullToEmpty(ex.getMessage()));
        span.tag(TraceSpanConstant.EXCEPTION_STACK_TRACE, StrUtil.nullToEmpty(ExceptionUtil.stacktraceToString(ex, 10240)));
    }

    /**
     * 填充tags数据
     */
    private void appendMdcInfo(MutableSpan span) {
        Map<String, String> mdcMap = MDC.getCopyOfContextMap();
        if (mdcMap == null || mdcMap.size() == 0) {
            return;
        }
        for (Map.Entry<String, String> e : mdcMap.entrySet()) {
            if (!IGNORE_KEYS.contains(e.getKey())) {
                span.tag(e.getKey(), e.getValue());
            }
        }

    }


    /**
     * 统计日志长度
     */
    private void countContentLength(MutableSpan span) {
        int tagCount = span.tagCount();
        int paramLength = this.calcParamLength(
                span.tag(TraceSpanConstant.HTTP_PARAM),
                span.tag(TraceSpanConstant.HTTP_HEADER),
                span.tag(TraceSpanConstant.RPC_PARAM),
                span.tag(TraceSpanConstant.SQL_GRAMMAR),
                span.tag(TraceSpanConstant.SQL_PARAM)
        );
        span.tag(TraceSpanConstant.LOG_HEIGHT, String.valueOf(tagCount));
        span.tag(TraceSpanConstant.LOG_LENGTH, String.valueOf(paramLength));
    }

    /**
     * 计算参数长度
     */
    private int calcParamLength(String... params) {
        if (params == null) {
            return 0;
        }
        int paramLength = 0;
        for (String param : params) {
            if (param == null) {
                continue;
            }
            paramLength += param.length();
        }
        return paramLength;
    }
}
